{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "YYZqhjGhK8nJ"
      },
      "source": [
        "##### Copyright 2018 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "V1-OvloqK4CX"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\"); { display-mode: \"form\" }\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "VKRNwMTxmZXx"
      },
      "source": [
        "# Concrete functions\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/guide/concrete_function\">\n",
        "    <img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />\n",
        "    View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/guide/concrete_function.ipynb\">\n",
        "    <img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />\n",
        "    Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/guide/concrete_function.ipynb\">\n",
        "    <img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />\n",
        "    View source on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/guide/concrete_function.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "VgQvKIvXockG"
      },
      "source": [
        "In the guide to [AutoGraph and `tf.functions`](function.ipynb) you saw how to use `tf.function`. This guide dives into the details of: \n",
        "\n",
        "* `tf.function` Tracing\n",
        "* `tf.function` Signatures\n",
        "* The Concrete functions generated by tracing:\n",
        "  * How to access them\n",
        "  * How to use them\n",
        "\n",
        "These details only become important:\n",
        "\n",
        "* If you're experiencing performance issues due to undesired tracing\n",
        "  of a `tf.funcion`.\n",
        "* When you need precise control over the TensorFlow Graphs generated by\n",
        "  `tf.function`. For example for exporting the model to\n",
        "  [TensorFlow Lite](https://tensorflow.org/lite/) using\n",
        "  `tf.lite.Converter.from_concrete_functions`.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "LwLkhiXIrT5m"
      },
      "source": [
        "## Background\n",
        "\n",
        "In TensorFlow 2, eager execution is on by default. TensorFlow's eager\n",
        "execution is an imperative programming environment that evaluates operations\n",
        "immediately, without building graphs. Operations return values instead\n",
        "of constructing a computational graph to run later. Here is a [detailed guide on eager\n",
        "execution](eager.ipynb).\n",
        "\n",
        "Running imperatively makes development and debugging\n",
        "more interactive, but doesn't allow for easy exporting.\n",
        "\n",
        "The `tf.function` API makes it possible to save models as graphs."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "GgJHWoOPrWtK"
      },
      "source": [
        "## Terminology\n",
        "\n",
        "The following terminology is used in this document:\n",
        "\n",
        "*   **Signature** - A description of the inputs and outputs for a set of operations.\n",
        "* **Polymorphic function** - Python callable that encapsulates several\n",
        "    concrete function graphs behind one API.  \n",
        "* **Concrete function** - Graph with a single signature.\n",
        "  \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "vq126xVZr2CL"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "3_jlx1_ir1hb"
      },
      "outputs": [],
      "source": [
        "import traceback\n",
        "import textwrap\n",
        "\n",
        "!pip install tf_nightly\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "AB6fMBGSofmm"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "TpkKbIBerbWi"
      },
      "source": [
        "## Create a `tf.function`\n",
        "\n",
        "Annotating a function with `tf.function` generates a *polymorphic function*\n",
        "containing those operations. All operations that are not annotated with\n",
        "`tf.function` will be evaluated with eager execution. The examples below show a quick example of `tf.function` usage."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "fPxeVDQ2mZXy"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def square(x):\n",
        "  return x*x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "yDO0s_LWBc9S"
      },
      "outputs": [],
      "source": [
        "square(2).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "8AhtB2hAsByK"
      },
      "source": [
        "Remember that the python decorator syntax just calls the decorator with the decorated object as input:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "p-KckIFcmZX2"
      },
      "outputs": [],
      "source": [
        "def pow(x,y):\n",
        "  return x ** y\n",
        "\n",
        "pow = tf.function(pow)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "1moLFTi2sToR"
      },
      "outputs": [],
      "source": [
        "pow(3,4).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "hMhMhl3lmZX6"
      },
      "source": [
        "### Attach a `tf.function` method to a `tf.Module`\n",
        "\n",
        "The `tf.function` can be optionally stored as part of a `tf.Module` object. The `tf.Module` class provides features for tracking variables and saving [checkpoints](checkpoints.ipynb) and [models](saved_model.ipynb).\n",
        "\n",
        "Classes like `keras.layers.Layer` and `keras.Model` are subclasses of Module."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ht2K2vpDmZX7"
      },
      "outputs": [],
      "source": [
        "class Pow(tf.Module):\n",
        "  def __init__(self, exponent):\n",
        "    self.exponent = tf.Variable(exponent, dtype = tf.float32, name='Pow/exponent')\n",
        "\n",
        "  @tf.function\n",
        "  def __call__(self, x):\n",
        "    return x ** self.exponent"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "1GVBiQfx8klb"
      },
      "outputs": [],
      "source": [
        "pow = Pow(3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "KqnKlZe58nPe"
      },
      "outputs": [],
      "source": [
        "pow.variables"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "RQBVvuxLtGz2"
      },
      "outputs": [],
      "source": [
        "pow(tf.constant(2.0)).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "kKiJrZlS5eeT"
      },
      "outputs": [],
      "source": [
        "pow.exponent.assign(4)\n",
        "pow(tf.constant(2.0)).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "t_HlUmTe5reS"
      },
      "outputs": [],
      "source": [
        "tf.saved_model.save(pow, 'pow')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "abpBoMRR5ypu"
      },
      "outputs": [],
      "source": [
        "reloaded_pow = tf.saved_model.load('pow')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "GA_a3CH956MR"
      },
      "outputs": [],
      "source": [
        "reloaded_pow(tf.constant(3.0)).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "rggP9vCK60R4"
      },
      "source": [
        "### Assign a `tf.function` as an attribute\n",
        "If you assign a `tf.Module` or a `tf.function` as an attribute of a module it will be serialized as well:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "kR2-Mwyg7GyO"
      },
      "outputs": [],
      "source": [
        "mod = tf.Module()\n",
        "mod.increment_by = tf.Variable(2.0)\n",
        "\n",
        "@tf.function\n",
        "def increment(x):\n",
        "  return x+mod.increment_by\n",
        "\n",
        "mod.inc = increment\n",
        "mod.inc(tf.constant(1.0)).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "rR5D4q7Z76Uo"
      },
      "outputs": [],
      "source": [
        "mod.cube = Pow(3)\n",
        "mod.cube(tf.constant(2.0)).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "vNIzKT8I8XzC"
      },
      "outputs": [],
      "source": [
        "mod.variables"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "OQqnUgjz7QOW"
      },
      "outputs": [],
      "source": [
        "tf.saved_model.save(mod, 'mod')\n",
        "reloaded_mod = tf.saved_model.load('mod')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Uvikzz7x7jB1"
      },
      "outputs": [],
      "source": [
        "reloaded_mod.inc(4.0).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "n7eVyFhT8Bq6"
      },
      "outputs": [],
      "source": [
        "reloaded_mod.cube(4.0).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Sj7P9wBFmZYW"
      },
      "source": [
        "### Interoperability with `tf.keras`\n",
        "\n",
        "Keras classes like `keras.Model` and `keras.layers.Layer` are fully compatible with `tf.function` and `tf.Module`.\n",
        "\n",
        "For example, build a simple model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "WYHD6NMWmZYX"
      },
      "outputs": [],
      "source": [
        "linear = tf.keras.Sequential([tf.keras.layers.Dense(units=1, input_shape=[1])])\n",
        "linear.compile(optimizer='adam', loss='mean_squared_error')\n",
        "linear.fit(x=[-1, 0, 1, 2, 3, 4], y=[-3, -1, 1, 3, 5, 7], epochs=50, verbose=0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "5rumHqek5oMN"
      },
      "outputs": [],
      "source": [
        "linear(tf.constant([[1],[2]]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "dx2HIyx_UI9N"
      },
      "source": [
        "Inspect it's variables"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "vLpavpL-86kn"
      },
      "outputs": [],
      "source": [
        "linear.variables"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "JEGvkaGH5hIZ"
      },
      "source": [
        "Now attach it to a `tf.Module`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "muWIMr6-5B8A"
      },
      "outputs": [],
      "source": [
        "module = tf.Module()\n",
        "module.linear = linear"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "SF24uHF59kH4"
      },
      "source": [
        "The `tf.Module` also tracks the `tf.Variable`s:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Z6ovPDgT5RYz"
      },
      "outputs": [],
      "source": [
        "module.variables"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "-ZyWZdJVtDp1"
      },
      "source": [
        "The `tf.Module` will export the contents of the `keras.Model` as well:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ZgNSWOoi6ioK"
      },
      "outputs": [],
      "source": [
        "tf.saved_model.save(module,'module')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "RS3j67MK9zVR"
      },
      "outputs": [],
      "source": [
        "reloaded = tf.saved_model.load('module')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "pxDCTrZd965w"
      },
      "outputs": [],
      "source": [
        "reloaded.linear([[1.0]])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ROhab4u7idvG"
      },
      "source": [
        "## Tracing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "zQaBkCHB-W0Q"
      },
      "source": [
        "The objects returned from `tf.function` are polymorphic functions. They will accept python objects, or `tf.Tensors` with any shape or `tf.dtype` as input.\n",
        "\n",
        "In the background TensorFlow builds `tf.Graph`s representing the calculation. \n",
        "This graph is wrapped in a python callable: a concrete function. Each concrete function can only handle a single input signature.\n",
        "\n",
        "`tf.function` traces the python function each time in needs to create a concrete function. The easiest way to see when a function is traced is to add a call to print: \n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "FgvSApDrmJGq"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def mul(a, b):\n",
        "  print('Tracing:\\n    {a}\\n    {b}\\n'.format(a=a, b=b))\n",
        "  return a*b"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "jxzhLeKV_chb"
      },
      "source": [
        "### Dtypes and shapes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "i_OnPSdOCI7J"
      },
      "source": [
        "If you call the polymorphic function with two different types of input, it will trace once for each:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "A2uKXbRMmTGj"
      },
      "outputs": [],
      "source": [
        "# Trace with ints\n",
        "mul(tf.constant(2), tf.constant(3)).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "a_rcmkWxCVwE"
      },
      "outputs": [],
      "source": [
        "# Trace with floats\n",
        "mul(tf.constant(2.0), tf.constant(3.0)).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Jmqq9rrDCdI4"
      },
      "source": [
        "When you call it again with the same input types, it dispatches to an existing function instead of tracing:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "KFtAlPZPqM7A"
      },
      "outputs": [],
      "source": [
        "# Call with ints again => no trace\n",
        "mul(tf.constant(10), tf.constant(10))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "mH920uLTCwK1"
      },
      "source": [
        "Changing the sizes of the inputs also triggers a trace (setting `tf.function(experimental_relax_shapes=True)` may reduce this): "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "eb_UByibqGny"
      },
      "outputs": [],
      "source": [
        "# Trace with vectors\n",
        "mul(tf.constant([1.0,3.0]), tf.constant(3.0)).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "eXm0fgVzqVxT"
      },
      "outputs": [],
      "source": [
        "# Trace with different-sized vectors\n",
        "mul(tf.constant([1.0,2.0,3.0, 4.0]), tf.constant(3.0))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "D-LoC287Ax5k"
      },
      "source": [
        "### Immutable python objects"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "RpkfdNeKFbiU"
      },
      "source": [
        "If you pass an immutable python object, like a `int`, `str`, or `tuple` to a `tf.function`, it executes a trace for each *value* of those python objects.\n",
        "\n",
        "This is useful to control what gets included in the `tf.Graph` (See: [The Autograph Guide](function.ipynb) for more details).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "-G1N7qh4BAhe"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def mul(a, b):\n",
        "  print('Tracing:\\n    {a}\\n    {b}\\n'.format(a=a, b=b))\n",
        "  return a*b"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ftCJz7HfDGof"
      },
      "outputs": [],
      "source": [
        "# Trace for a=3.0\n",
        "mul(3.0, tf.constant(3.0)).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "C8hiJCP_Du0I"
      },
      "outputs": [],
      "source": [
        "# Don't trace for a=3.0 the second time:\n",
        "mul(3.0, tf.constant(3.0)).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Sm8Cq0uTt88j"
      },
      "source": [
        "Caution: It is easy to cause many traces by passing unique python values. This can be a significant performance problem. Often passing a `tf.Tensor` value is the solution.\n",
        "\n",
        "This loop traces the function for each unique int:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "4LTTGqQryjsj"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def power(a,b):\n",
        "  print('Tracing \"power\": a={}'.format(a))\n",
        "  return a**b"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "yyGenHr_t8LK"
      },
      "outputs": [],
      "source": [
        "p = tf.constant(2)\n",
        "for n in range(12):\n",
        "  power(n,p)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "xYgJ9l7_yY20"
      },
      "source": [
        "On the second run each int has been traced, so there's no tracing to do:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "duPVVOg3yUCl"
      },
      "outputs": [],
      "source": [
        "p = tf.constant(2)\n",
        "for n in range(12):\n",
        "  power(n,p)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "6bX9b1t0uTQX"
      },
      "source": [
        "To avoid excess retracing be sure to pass a `tf.Tensor` instead of python numbers or strings:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "F78ik9psuPn7"
      },
      "outputs": [],
      "source": [
        "p = tf.constant(2)\n",
        "for n in tf.range(12):\n",
        "  power(n,p)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "N-wwjPtF1Q25"
      },
      "source": [
        "To shut off tracing altogether, pass a signature to the `tf.function` decorator:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "9EcEZfZY1Ylg"
      },
      "outputs": [],
      "source": [
        "@tf.function(input_signature=(\n",
        "    tf.TensorSpec(shape=[], dtype=tf.float32),\n",
        "    tf.TensorSpec(shape=[], dtype=tf.float32),)\n",
        ")\n",
        "def power_with_sig(a,b):\n",
        "  print('Tracing \"power_with_sig\"')\n",
        "  return a**b"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Pff46f4U1-6R"
      },
      "outputs": [],
      "source": [
        "power_with_sig(3.0, 3.0).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "cBT2IFiZ1wLS"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "  power_with_sig(tf.constant([1.0,2.0,3.0]),tf.constant(3.0))\n",
        "  assert False\n",
        "except ValueError:\n",
        "  traceback.print_exc(limit=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "kjeEQ4J4tt81"
      },
      "source": [
        "### Example: Dropout"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "-QbejAfEDLNX"
      },
      "source": [
        "Retracing for specific values gives you control over what code gets generated by the `tf.function`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "kIgq7vO1sXh7"
      },
      "outputs": [],
      "source": [
        "class Dropout(tf.Module):\n",
        "  def __init__(self, rate, name=None):\n",
        "    super(Dropout, self).__init__(name)\n",
        "    self.rate = tf.Variable(rate, dtype = tf.float32, trainable=False)\n",
        "\n",
        "  @tf.function\n",
        "  def __call__(self, x, training=True):\n",
        "    print(textwrap.dedent(\"\"\"\n",
        "                          Tracing \"Dropout\":\n",
        "                              training = {}\n",
        "                              x = {}\n",
        "                              name = {:s}\n",
        "                          \"\"\".format(training, x, self.name)))\n",
        "    if training:\n",
        "      print('    - Train branch\\n')\n",
        "      mask = tf.random.uniform(x.shape) > self.rate\n",
        "      return x * tf.cast(mask, tf.float32)/self.rate\n",
        "    else:\n",
        "      print('    - Test branch\\n')\n",
        "      return x"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "o7TFSeINqhtc"
      },
      "source": [
        "Create an instance of this simple `Dropout` layer:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "5DHE1io1pN2t"
      },
      "outputs": [],
      "source": [
        "dropout = Dropout(0.5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "CE4fuvUP0BWE"
      },
      "source": [
        "The first time you call it with a python `training=True` as input, it traces the `training` branch:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "O5xIiGATqNQY"
      },
      "outputs": [],
      "source": [
        "dropout(tf.range(10, dtype=tf.float32), training=True).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Cup3pXHI0WWH"
      },
      "source": [
        "The second time, it doesn't need to re-trace the branch:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "VaMDgyM8zpge"
      },
      "outputs": [],
      "source": [
        "dropout(tf.range(10, dtype=tf.float32), training=True).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "yhiyPC7s0dF9"
      },
      "source": [
        "Passing `training=False` triggers a trace on the first run since this is a different python value:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "mgaR5PFTzk9A"
      },
      "outputs": [],
      "source": [
        "dropout(tf.range(10, dtype=tf.float32), training=False).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "DVeAfMkbzn-V"
      },
      "outputs": [],
      "source": [
        "dropout(tf.range(10, dtype=tf.float32), training=False).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "G6b57DUG32b7"
      },
      "source": [
        "If you pass a `bool` tensor, it uses TensorFlow autograph rewrite the `if` to a `tf.cond`m and traces both branches:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "s_FsCDEaqYqq"
      },
      "outputs": [],
      "source": [
        "dropout(tf.range(10, dtype=tf.float32), training=tf.constant(False)).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "UbmPUjhI8Rgf"
      },
      "source": [
        "This captures the control flow in a single concrete function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "BViHlbCN0wss"
      },
      "outputs": [],
      "source": [
        " dropout(tf.range(10, dtype=tf.float32), training=tf.constant(True)).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "p8eT9zKXp8F5"
      },
      "outputs": [],
      "source": [
        "dropout(tf.range(10, dtype=tf.float32), training=tf.constant(False)).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "smyOmjl2GCdt"
      },
      "source": [
        "### Other python objects\n",
        "\n",
        "Since the generated `tf.Graphs` cannot contain complex python objects, these are included by tracing and variable capture. \n",
        "\n",
        "The `tf.function` runs a separate trace for each **instance**. So each trace includes its own variables, and can set its behavior based on the instance.\n",
        "\n",
        "The most common usage is on methods of Module, Layer or Module:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "C6SBvOEdO5Qw"
      },
      "outputs": [],
      "source": [
        "dropout_a = Dropout(0.5, name='dropout_a')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "uOC8COMxQY0J"
      },
      "outputs": [],
      "source": [
        "print(dropout_a(tf.range(10, dtype=tf.float32), True).numpy())\n",
        "print(dropout_a(tf.range(10, dtype=tf.float32), True).numpy())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "mISom9-ZTf5g"
      },
      "outputs": [],
      "source": [
        "dropout_b = Dropout(0.5, name='dropout_b')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "95OE0xxeQ4t0"
      },
      "outputs": [],
      "source": [
        "print(dropout_b(tf.range(10, dtype=tf.float32), True).numpy())\n",
        "print(dropout_b(tf.range(10, dtype=tf.float32), True).numpy())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Ts8nQVMcQfsK"
      },
      "source": [
        "But the behavior is the same on a stand-alone `tf.function`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "0xVE3zoLs61l"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def run(callable, x):\n",
        "  print('Tracing \"run\":\\n    callable = {}\\n    x = {}\\n'.format(callable, x))\n",
        "  return callable(x)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "WAs-suN2tjH0"
      },
      "outputs": [],
      "source": [
        "def plus_1(x):\n",
        "  return x+1\n",
        "\n",
        "print(run(plus_1, tf.constant(2.0)).numpy())\n",
        "print(run(plus_1, tf.constant(5.0)).numpy())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "g_m9M64oXFeW"
      },
      "source": [
        "The tracing one `tf.function` can trigger tracing in another:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "lnib6kNTu4bF"
      },
      "outputs": [],
      "source": [
        "print(run(dropout, tf.range(10.0)).numpy())\n",
        "print(run(dropout, tf.range(10.0)).numpy())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "4DE76z8KiVTG"
      },
      "source": [
        "### Weak references"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "LKaSfW5likUT"
      },
      "source": [
        "Caution: Each trace only keeps a [weak-reference](https://docs.python.org/3/library/weakref.html) to any `tf.Variable`. If the variable is not kept alive by another reference, the trace may become unusable."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "yGUY0lVD8nMT"
      },
      "source": [
        "For example here's a `tf.function` that refers to `var` from the enclosing scope:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "zwwnJwXR8PEb"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def plus_var(x):\n",
        "  print('Tracing \"plus_var\":\\n    x = {}\\n    var = {}\\n\\n'.format(x, var.name))\n",
        "  return x + var"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "9qmxAmFl8v5p"
      },
      "source": [
        "Trace the function with one variable:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Nebd4e5y8UYa"
      },
      "outputs": [],
      "source": [
        "var = tf.Variable(1, name=\"IntVar\")\n",
        "plus_var(tf.constant([1,2])).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "DaqblviJ8y6v"
      },
      "source": [
        "And with another variable:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "s3EwG-yt8Zgx"
      },
      "outputs": [],
      "source": [
        "var = tf.Variable(2.0, name=\"FloatVar\")\n",
        "plus_var(tf.constant([2.0, 10.0])).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "cJL43oZN86o4"
      },
      "source": [
        "That worked, but because you no longer have a reference to `\"IntVar\"`, that first trace is broken:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Yb4RkESA8gs2"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "  plus_var(tf.constant([1,2])).numpy()\n",
        "  assert False\n",
        "except tf.errors.FailedPreconditionError:\n",
        "  traceback.print_exc(limit=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "235EbM2HmZYB"
      },
      "source": [
        "## Accessing concrete function\n",
        "\n",
        "In the previous section you saw the conditions for triggering a new trace of a polymorphic `tf.function`. Each trace generates a new concrete function.\n",
        "\n",
        "When you save `tf.Module` as a `tf.saved_model` It's those concrete functions that define the `tf.Graph`s that are exported. You don't save a `tf.function` you save the concrete functions that are created by tracing. \n",
        "\n",
        "To get a concrete function from the polymorphic `tf.function` you need to define the signature. Either:\n",
        "\n",
        "*   Pass an `input_signature` to `tf.function`, and call the \n",
        "    `get_concrete_function()` method.\n",
        "*   Pass a list of `tf.TensorSpec`s to `get_concrete_function`: `tf.TensorSpec(shape=[1], dtype=tf.float32)`.\n",
        "*   Pass an example tensor of the correct shape and type to\n",
        "    `get_concrete_function`: `tf.constant(1., shape=[1])`.\n",
        "\n",
        "The following example shows how to define the `input_signature` parameter for\n",
        "`tf.function`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "O-tE5Pn6h1DU"
      },
      "source": [
        "#### Using `input_signature`"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "hsMR5JBdhykq"
      },
      "source": [
        "Specify input tensors in the call to `tf.function` as shown below.\n",
        "This `tf.function`can only execute on tensors that match the specified signatutre.\n",
        "\n",
        "A `None` in the `shape` acts a wildcard. So this these `tf.TensroSpec` say \"A float32 vector of any length\".\n",
        "\n",
        "This pattern can be very important if your `tf.function` is expected to handle sequences of different length, or images of different sizes for each batch (See [Transformer](../tutorials/text/transformer.ipynb) and [Deep Dream](../tutorials/generative/deepdream.ipynb) tutrorials for example)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "INJhW9c5htAK"
      },
      "outputs": [],
      "source": [
        "@tf.function(input_signature=(\n",
        "    tf.TensorSpec(shape=[None], dtype=tf.float32),\n",
        "    tf.TensorSpec(shape=[None], dtype=tf.float32),)\n",
        ")\n",
        "def power_with_sig(a,b):\n",
        "  print('Tracing \"power_with_sig\"\\n')\n",
        "  return a**b"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Q0U5ytANi4V7"
      },
      "source": [
        "Calling `get_concrete_function` will execute the trace (if necessary), and return a concrete function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "yR3tqv1pijvh"
      },
      "outputs": [],
      "source": [
        "p = power_with_sig.get_concrete_function()\n",
        "type(p)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "XvGhi41EipeQ"
      },
      "outputs": [],
      "source": [
        "p(tf.constant([2.0,3.0,4.0]), tf.constant([5.0,4.0,3.0])).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "bBKfDEeMz43m"
      },
      "source": [
        "### Using `get_concrete_function`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "IGlmMUSK0c3H"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def power(a,b):\n",
        "  print('Tracing \"power\"\\n')\n",
        "  return a**b"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "NYbxGotFqoDB"
      },
      "outputs": [],
      "source": [
        "float_power = power.get_concrete_function(\n",
        "  a = tf.TensorSpec(shape=[], dtype=tf.float32),\n",
        "  b = tf.TensorSpec(shape=[], dtype=tf.float32))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "1cXOt1g0qpCL"
      },
      "outputs": [],
      "source": [
        "float_power(tf.constant(3.0),tf.constant(3.0))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "0dhioJFVbyKJ"
      },
      "source": [
        "Remember that you can also pass tensors to `get_concrete_function`, in that case it returns the concrete function that would run for those inputs:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "dzeSYHN1cGke"
      },
      "outputs": [],
      "source": [
        "row = tf.range(10)\n",
        "col = tf.constant([[1],[2],[3]])\n",
        "\n",
        "concrete_power = power.get_concrete_function(a = row, b = col)\n",
        "concrete_power(row, col).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "JcsoDz8g9nJT"
      },
      "source": [
        "## Using a concrete function"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "glxZkvRly1Am"
      },
      "source": [
        "A concrete function only accepts tensors as input:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "1AIbwzSz0aCj"
      },
      "outputs": [],
      "source": [
        "float_power(tf.constant(2.0), tf.constant(3.0)).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "7xq-82FokX_d"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "  float_power(2.0,3.0)\n",
        "  assert False\n",
        "except (ValueError, TypeError):\n",
        "  traceback.print_exc(limit=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "V1eBDquxy5m5"
      },
      "source": [
        "It also only accepts inputs of the correct dtype:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "nSdGa_IJkayV"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "  float_power(tf.constant(1),tf.constant(3))\n",
        "  assert False\n",
        "except tf.errors.InvalidArgumentError:\n",
        "  traceback.print_exc(limit=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "kRc4YAxzsSdM"
      },
      "source": [
        "But it will try to execute even if the input tensors do not match the expected shape:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "MT-Z0cwrrrgZ"
      },
      "outputs": [],
      "source": [
        "float_power(tf.constant([1.,2.,3.,4.,5.]),tf.constant(3.)).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "zWr690sEqwu0"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "  float_power(tf.constant([1.,2.,3.]),tf.constant([4., 5.])).numpy()\n",
        "  assert False\n",
        "except tf.errors.InvalidArgumentError:  \n",
        "  traceback.print_exc(limit=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "PtHTfAtamxcu"
      },
      "source": [
        "By inspecting the concrete function you can see its inputs and outputs:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "d-8tBcOHmw7x"
      },
      "outputs": [],
      "source": [
        "print(float_power.structured_input_signature)\n",
        "print(float_power.structured_outputs)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "64z5eSmMzv8l"
      },
      "source": [
        "## Python Objects in signatures"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "1RTj0n2R0jGH"
      },
      "source": [
        "As you saw when tracing, each python object generates a new trace. Concrete functions represent a single `tf.Graph`, they don't do any retracing. When you call `get_concrete_function` with a python object as one of the arguments the object is **bound** to the function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "HxQtkTRYz-8u"
      },
      "outputs": [],
      "source": [
        "cube = power.get_concrete_function(\n",
        "    a = tf.TensorSpec([], dtype=tf.float32),\n",
        "    b = 3.0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Iu8dRqCQ0SGW"
      },
      "source": [
        "This `cube` function no longer has a `b` argument:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "duRJIrDU0L-I"
      },
      "outputs": [],
      "source": [
        "print(cube.structured_input_signature)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "PnEQtxxA3OgX"
      },
      "outputs": [],
      "source": [
        "cube(tf.constant(10.0)).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "hL76DCMVsdet"
      },
      "source": [
        "This is very similar to the way that standard python classes bind methods, and applies equally when you run `get_concrete_function` from a method:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "dV2R0AOIsdJs"
      },
      "outputs": [],
      "source": [
        "class Greeter(object):\n",
        "  def __init__(self, greeting):\n",
        "    self.greeting = greeting\n",
        "\n",
        "  def greet(self, who):\n",
        "    return \" \".join([self.greeting, who])\n",
        "\n",
        "p = Greeter(\"Hello\")\n",
        "m = p.greet\n",
        "print(m)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "j9sEWo-wtP_K"
      },
      "outputs": [],
      "source": [
        "print(m(\"TensorFlow!\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "uv4X41_Htqof"
      },
      "source": [
        "When you have a `tf.function` decorating a method, similar rules apply:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "h131HNsj3nYm"
      },
      "outputs": [],
      "source": [
        "class MyModel(tf.Module):\n",
        "  def __init__(self, ins, outs):\n",
        "    initializer = tf.initializers.GlorotNormal()\n",
        "    self.W = tf.Variable(initializer([ins, outs]))\n",
        "    self.B = tf.Variable(tf.zeros([outs], dtype = tf.float32))\n",
        "\n",
        "  @tf.function\n",
        "  def run(self, x):\n",
        "    print('Tracing \"MyModule\":\\n    x={}\\n'.format(x))\n",
        "    return tf.matmul(x, self.W)+self.B"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "nGtPiUpZ3y0j"
      },
      "outputs": [],
      "source": [
        "mod = MyModel(ins=5, outs=3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "3OVnEqJY6BHU"
      },
      "outputs": [],
      "source": [
        "mod.run([[1.0,1.0,1.0, 1.0, 1.0]]).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ca4tB9QguPBg"
      },
      "source": [
        "If you call the method's `.get_concrete_function`, the `self` is automatically bound as the first argument:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "0buNZlUw6XNN"
      },
      "outputs": [],
      "source": [
        "concrete_run = mod.run.get_concrete_function(x = tf.TensorSpec([None, None]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "D-v1YUN562wH"
      },
      "outputs": [],
      "source": [
        "concrete_run(tf.constant([[1.0,1.0,1.0, 1.0, 1.0],\n",
        "                          [2.0,2.0,2.0, 2.0, 2.0]])).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "voRoyyNBuq5-"
      },
      "source": [
        "See how `self` is no longer part of the input signature:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "t4V1ChR97uNS"
      },
      "outputs": [],
      "source": [
        "print(concrete_run.structured_input_signature)\n",
        "print(concrete_run.structured_outputs)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "YYHxaxplBaSt"
      },
      "source": [
        "## Changes for TensorFlow 2.3\n",
        "\n",
        "The following changes are currently available in TensorFlow nightly, and will be available in TensorFlow 2.3."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "MTk5OR1sDLU9"
      },
      "source": [
        "### Inspecting concrete function signatures\n",
        "\n",
        "Printing a `ConcreteFunction` displays a summary of its input arguments (with types) and its output type:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "NxQOvjDFD52f"
      },
      "outputs": [],
      "source": [
        "print(float_power)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "EjgE_2TQD8oh"
      },
      "source": [
        "For polymorphic functions, the `pretty_printed_concrete_signatures()` method can be used to display a list of all signatures that have been traced so far:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "YwQTaUnJDSqP"
      },
      "outputs": [],
      "source": [
        "print(power.pretty_printed_concrete_signatures())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Q0E84pJUEI01"
      },
      "source": [
        "### Bound arguments\n",
        "\n",
        "As discussed above, when a concrete function is created, any arguments that are set to non-Tensor values are *bound* to those values.  Prior to TensorFlow 2.3, these bound arguments were simply removed from the concrete function's signature.  Starting with TensorFlow 2.3, they remain present in the signature, but have their default value set to the bound value.  For example, the argument `b` has its default value set to `2` when we create the following concrete function:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "FDIiQ2BqFvdQ"
      },
      "outputs": [],
      "source": [
        "square = power.get_concrete_function(a=tf.TensorSpec(None, tf.float32), b=2)\n",
        "print(square)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "7u_Zl_WyF5vu"
      },
      "source": [
        "When calling `square`, you may either omit `b`, or explicitly set it to the value that was used during tracing:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "NVdhWX3MGAdo"
      },
      "outputs": [],
      "source": [
        "print(square(tf.constant(5.0)))  # ok: use default value\n",
        "print(square(tf.constant(5.0), b=2))  # ok: explicit value == bound value"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "RTWRfpO8GJre"
      },
      "source": [
        "Setting a bound argument to any value other than its default is an error:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "AvIkTXXRGOhN"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "  print(square(tf.constant(10.0), b=4.0).numpy())  # error: explicit value != bound value\n",
        "  assert False\n",
        "except TypeError:\n",
        "  traceback.print_exc(limit=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "RTZYDSnQCMOx"
      },
      "source": [
        "### CompositeTensor arguments\n",
        "\n",
        "Starting with TensorFlow 2.3, you can define concrete functions that accept `CompositeTensor`s (such as `tf.RaggedTensor`, `tf.Dataset` iterators, and `tf.SparseTensor`).  E.g.:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "CbfBbUq4Gmwp"
      },
      "outputs": [],
      "source": [
        "ragged_power = power.get_concrete_function(\n",
        "    a = tf.RaggedTensorSpec([None, None], dtype=tf.float32),\n",
        "    b = tf.TensorSpec([], dtype=tf.float32)\n",
        ")\n",
        "print(ragged_power(tf.ragged.constant([[1.0, 2.0], [3.0]]),\n",
        "                   tf.constant(2.0)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "U2E_PAQeG4R-"
      },
      "source": [
        "### Nested arguments\n",
        "\n",
        "Starting with TensorFlow 2.3, you can define concrete functions that accept nested `dict`s, `list`s, `tuple`s, `namedtuple`s, and [`attr`](https://www.attrs.org/en/stable/)s of tensors (or composite tensors).  E.g.:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "r7SnsGIfHYWt"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def sum_features(features):\n",
        "  assert isinstance(features, dict)\n",
        "  return sum(t for t in features.values())\n",
        "\n",
        "features = {\n",
        "    'a': tf.constant([[1.0], [2.0], [8.0]]),\n",
        "    'b': tf.constant([[5.0], [6.0], [0.0]]),\n",
        "    'c': tf.ragged.constant([[10.0], [20.0, 30.0, 40.0], [50.0]])\n",
        "}\n",
        "print(sum_features(features))\n",
        "\n",
        "concrete_sum_features = sum_features.get_concrete_function(features)\n",
        "print(concrete_sum_features(features))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "9vnmuNhoH-aV"
      },
      "source": [
        "The value passed to a concrete function must have the same nesting structure that was used when the concrete function was traced.  If you pass in a different structure, then a `TypeError` is raised, with information about the structure that was accepted and the structure that was used:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Oh0WF_N9IGPt"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "  concrete_sum_features({'a': tf.constant([[1.0]])})\n",
        "  assert False\n",
        "except TypeError:\n",
        "  traceback.print_exc(limit=1)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "concrete_function.ipynb",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
